import numpy as np
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.metrics import accuracy_score, silhouette_score, calinski_harabasz_score
from sklearn.preprocessing import StandardScaler
from scipy.optimize import linear_sum_assignment

# 加载鸢尾花数据集
print("loading")
iris = datasets.load_iris()
X = iris.data
true_labels = iris.target

# 数据标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# K-means聚类
def k_means_clustering(X, n_clusters=3, max_iters=100, tol=1e-4):
    print("kmeans")
    np.random.seed(42)
    centroids = X[np.random.choice(X.shape[0], n_clusters, replace=False)]

    for _ in range(max_iters):
        distances = np.sqrt(((X - centroids[:, np.newaxis]) ** 2).sum(axis=2))
        closest_centroids = np.argmin(distances, axis=0)

        new_centroids = np.array([X[closest_centroids == k].mean(axis=0) for k in range(n_clusters)])

        if np.allclose(centroids, new_centroids, atol=tol):
            break

        centroids = new_centroids

    labels = closest_centroids
    return labels, centroids

# DBSCAN聚类
def dbscan_clustering(X, eps=5, min_samples=10):
    print("dbscan")
    labels = np.full(X.shape[0], -1, dtype=int)
    label = 0

    for i in range(X.shape[0]):
        if labels[i] == -1:
            neighbors = np.where(np.linalg.norm(X[i] - X, axis=1) < eps)[0]
            if len(neighbors) < min_samples:
                labels[i] = -2  # Noise
            else:
                labels[i] = label
                for neighbor in neighbors:
                    if labels[neighbor] == -1:
                        labels[neighbor] = label
                    elif labels[neighbor] == -2:
                        labels[neighbor] = label
                label += 1

    return labels

# 评估函数
def evaluate_clustering(labels, true_labels):
    filtered_indices = labels != -2  # Remove noise
    labels = labels[filtered_indices]
    true_labels = true_labels[filtered_indices]
    if len(labels) == 0:
        return 0, 0, 0  # Avoid division by zero if all points are noise

    # Align labels using Hungarian algorithm
    unique_labels = np.unique(labels)
    unique_true_labels = np.unique(true_labels)
    cost_matrix = np.zeros((len(unique_labels), len(unique_true_labels)))
    for i, ul in enumerate(unique_labels):
        for j, utl in enumerate(unique_true_labels):
            cost_matrix[i, j] = np.sum((labels == ul) & (true_labels != utl))
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    aligned_labels = np.zeros_like(labels)
    for i, ul in enumerate(unique_labels):
        if i < len(col_ind):  # Ensure we do not exceed the number of true labels
            aligned_labels[labels == ul] = unique_true_labels[col_ind[i]]

    accuracy = accuracy_score(true_labels, aligned_labels)
    silhouette = silhouette_score(X_scaled[filtered_indices], labels)
    calinski_harabasz = calinski_harabasz_score(X_scaled[filtered_indices], labels)
    return accuracy, silhouette, calinski_harabasz

# 运行K-means
kmeans_labels, kmeans_centroids = k_means_clustering(X_scaled)
kmeans_accuracy, kmeans_silhouette, kmeans_calinski_harabasz = evaluate_clustering(kmeans_labels, true_labels)

# 运行DBSCAN
dbscan_labels = dbscan_clustering(X_scaled, eps=2.5, min_samples=15)
dbscan_accuracy, dbscan_silhouette, dbscan_calinski_harabasz = evaluate_clustering(dbscan_labels, true_labels)

print("K-means:")
print(f"Accuracy: {kmeans_accuracy:.2f}")
print(f"Silhouette Score: {kmeans_silhouette:.2f}")
print(f"Calinski-Harabasz Index: {kmeans_calinski_harabasz:.2f}")

print("\nDBSCAN:")
print(f"Accuracy: {dbscan_accuracy:.2f}")
print(f"Silhouette Score: {dbscan_silhouette:.2f}")
print(f"Calinski-Harabasz Index: {dbscan_calinski_harabasz:.2f}")

# 可视化结果
plt.figure(figsize=(15, 5))

plt.subplot(1, 3, 1)
plt.scatter(X_scaled[:, 0], X_scaled[:, 1], c=kmeans_labels)
plt.title('K-means Clustering')

plt.subplot(1, 3, 2)
plt.scatter(X_scaled[:, 0], X_scaled[:, 1], c=dbscan_labels)
plt.title('DBSCAN Clustering')

plt.subplot(1, 3, 3)
plt.scatter(X_scaled[:, 0], X_scaled[:, 1], c=true_labels)
plt.title('True Clustering')

plt.show()